package com.sjw.workflow.agent;

import com.sjw.workflow.core.Agent;
import com.sjw.workflow.core.AgentInput;
import com.sjw.workflow.core.AgentResult;
import com.sjw.workflow.core.WorkflowContext;
import com.sjw.workflow.core.ExecutionStatus;
// import dev.langchain4j.model.chat.ChatLanguageModel;
// import dev.langchain4j.model.openai.OpenAiChatModel;
// import dev.langchain4j.service.AiServices;
// import dev.langchain4j.service.UserMessage;
// import org.slf4j.Logger;
// import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.HashMap;
import java.util.concurrent.CompletableFuture;
import java.util.function.Function;

/**
 * Agent基础实现类
 * 提供通用的Agent功能，包括LLM调用和工具使用
 */
public abstract class BaseAgent implements Agent {
    
    // private static final Logger logger = LoggerFactory.getLogger(BaseAgent.class);
    
    protected final String id;
    protected final String name;
    protected final String description;
    // protected ChatLanguageModel chatModel;
    // protected AgentAssistant assistant;
    
    public BaseAgent(String id, String name, String description) {
        this.id = id;
        this.name = name;
        this.description = description;
        initializeChatModel();
        initializeAssistant();
    }
    
    /**
     * 初始化聊天模型
     */
    private void initializeChatModel() {
        // 使用OpenAI模型，实际使用时需要配置API Key
        // this.chatModel = OpenAiChatModel.builder()
        //         .apiKey("demo-key") // 实际使用时需要真实的API Key
        //         .modelName("gpt-3.5-turbo")
        //         .temperature(0.7)
        //         .build();
    }
    
    /**
     * 初始化AI助手
     */
    private void initializeAssistant() {
        // this.assistant = AiServices.builder(AgentAssistant.class)
        //         .chatLanguageModel(chatModel)
        //         .build();
    }
    
    @Override
    public String getId() {
        return id;
    }
    
    @Override
    public String getName() {
        return name;
    }
    
    @Override
    public String getDescription() {
        return description;
    }
    
    @Override
    public boolean isReady(WorkflowContext context) {
        return true; // 默认总是准备就绪，子类可以重写
    }
    
    @Override
    public CompletableFuture<AgentResult> execute(AgentInput input, WorkflowContext context) {
        return CompletableFuture.supplyAsync(() -> {
            long startTime = System.currentTimeMillis();
            System.out.println("Starting execution of agent: " + id);
            
            try {
                // 预处理输入
                Map<String, Object> processedInput = preprocessInput(input, context);
                
                // 执行Agent特定逻辑
                Map<String, Object> output = executeInternal(processedInput, context);
                
                // 后处理输出
                Map<String, Object> processedOutput = postprocessOutput(output, context);
                
                long executionTime = System.currentTimeMillis() - startTime;
                System.out.println("Completed execution of agent: " + id + " in " + executionTime + "ms");
                
                AgentResult result = new AgentResult(id, ExecutionStatus.SUCCESS, processedOutput);
                result.setExecutionTime(executionTime);
                return result;
                
            } catch (Exception e) {
                long executionTime = System.currentTimeMillis() - startTime;
                System.err.println("Failed to execute agent: " + id + " - " + e.getMessage());
                
                AgentResult result = new AgentResult(id, ExecutionStatus.FAILED, new HashMap<>());
                result.setError(e.getMessage());
                result.setExecutionTime(executionTime);
                return result;
            }
        });
    }
    
    /**
     * 预处理输入数据
     */
    protected Map<String, Object> preprocessInput(AgentInput input, WorkflowContext context) {
        Map<String, Object> processed = new HashMap<>();
        
        if (input.getData() != null) {
            processed.putAll(input.getData());
        }
        
        // 添加上下文信息
        processed.put("_context", context);
        processed.put("_agentId", id);
        processed.put("_agentName", name);
        
        return processed;
    }
    
    /**
     * 后处理输出数据
     */
    protected Map<String, Object> postprocessOutput(Map<String, Object> output, WorkflowContext context) {
        Map<String, Object> processed = new HashMap<>();
        
        if (output != null) {
            processed.putAll(output);
        }
        
        // 添加执行元数据
        processed.put("_agentId", id);
        processed.put("_agentName", name);
        processed.put("_executionTime", System.currentTimeMillis());
        
        return processed;
    }
    
    /**
     * 子类需要实现的核心执行逻辑
     */
    protected abstract Map<String, Object> executeInternal(Map<String, Object> input, WorkflowContext context) throws Exception;
    
    /**
     * 调用LLM进行文本处理
     */
    protected String callLLM(String prompt) {
        try {
            // return assistant.processText(prompt);
            return "LLM processing: " + prompt; // 模拟LLM响应
        } catch (Exception e) {
            System.out.println("LLM call failed, using fallback response: " + e.getMessage());
            return "LLM processing failed: " + e.getMessage();
        }
    }
    
    /**
     * 调用工具函数
     */
    protected <T, R> R callTool(Function<T, R> tool, T input) {
        try {
            return tool.apply(input);
        } catch (Exception e) {
            System.err.println("Tool call failed: " + e.getMessage());
            throw new RuntimeException("Tool execution failed: " + e.getMessage(), e);
        }
    }
    
    /**
     * 获取上下文中的Agent结果
     */
    protected Map<String, Object> getAgentResult(WorkflowContext context, String agentId) {
        if (context.isAgentCompleted(agentId)) {
            return context.getAgentResult(agentId).getOutput();
        }
        return new HashMap<>();
    }
    
    /**
     * 获取上下文中的全局数据
     */
    protected Object getGlobalData(WorkflowContext context, String key) {
        return context.getGlobalData(key);
    }
    
    /**
     * 设置上下文中的全局数据
     */
    protected void setGlobalData(WorkflowContext context, String key, Object value) {
        context.setGlobalData(key, value);
    }
}

/**
 * AI助手接口，用于与LLM交互
 */
// interface AgentAssistant {
//     
//     @UserMessage("请处理以下文本：{{text}}")
//     String processText(String text);
// }
